import pandas as pd
import os
import re
import jieba
from sklearn.model_selection import train_test_split


def clean_text(text: str) -> str:
    """文本清洗函数"""
    # 替换多个空格为单个空格
    text = re.sub(r'\s+', ' ', text)

    # 去除多余标点
    text = re.sub(r'([!?,.:;。，？！：；])\1+', r'\1', text)

    # 去除各种特殊符号组合
    special_patterns = [
        '//:', '//', '///', '////',  # 斜杠组合
        '::', ':::', '::::',  # 冒号组合
        '：：', '：：：', '：：：：',  # 中文冒号组合
        '//：', '///：', '////：',  # 混合冒号组合
        '：//', '：///', '：////',  # 反向混合冒号组合
        '//：：', '///：：',  # 多重混合组合
        '：：//', '：：///'  # 反向多重混合组合
    ]
    for pattern in special_patterns:
        text = text.replace(pattern, '')

    # 统一常见表示形式
    text = text.replace('really', '非常')
    text = text.replace('very', '很')
    text = text.replace('so', '如此')

    # 去除URL
    text = re.sub(r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', '', text)

    # 去除@用户
    text = re.sub(r'@[\w\-]+', '', text)

    # 去除表情符号
    text = re.sub(r'\[.*?\]', '', text)

    # 去除数字
    text = re.sub(r'\d+', '', text)

    # 去除多余的空格
    text = re.sub(r'\s+', ' ', text)

    return text.strip()


def prepare_weibo_data():
    # 读取本地数据文件
    data_path = '../data/weibo_senti_100k.csv'

    print("正在读取数据...")
    df = pd.read_csv(data_path, encoding='utf-8')

    print("正在进行数据清洗...")
    # 对文本进行清洗
    df['review'] = df['review'].apply(clean_text)

    # 使用分层抽样划分训练集和测试集
    train_df, test_df = train_test_split(
        df,
        test_size=0.2,  # 20%用于测试
        random_state=42,  # 设置随机种子以确保结果可复现
        stratify=df['label']  # 使用标签进行分层抽样
    )

    # 保存处理后的数据
    train_df.to_csv('../data/weibo_train.csv', index=False, encoding='utf-8')
    test_df.to_csv('../data/weibo_test.csv', index=False, encoding='utf-8')

    print("数据集准备完成！")
    print(f"训练集大小: {len(train_df)} 条数据")
    print(f"测试集大小: {len(test_df)} 条数据")

    # 打印测试集中各类别的分布
    print("\n测试集类别分布:")
    print(test_df['label'].value_counts())


if __name__ == '__main__':
    prepare_weibo_data()